import pandas as pd
import numpy as np
import joblib
from data_preprocessing import feature_engineering

def predict_fraud(input_data):
    """
    使用训练好的模型预测单条数据

    参数:
        input_data (dict): 包含所有特征的字典

    返回:
        dict: 包含欺诈概率和预测结果(0或1)
    """
    # 加载预处理对象
    preprocessor = joblib.load('preprocessor.pkl')
    selector = joblib.load('feature_selector.pkl')
    label_encoders = joblib.load('label_encoders.pkl')
    model = joblib.load('best_model.pkl')

    # 转换为DataFrame
    input_df = pd.DataFrame([input_data])

    # 应用相同的特征工程
    input_df = feature_engineering(input_df)

    # 移除不需要的列
    input_df = input_df.drop(['policy_id', 'policy_bind_date', 'incident_date'], axis=1, errors='ignore')

    # 获取特征列表
    numeric_features = preprocessor.transformers_[0][2]
    categorical_features = list(label_encoders.keys())

    # 预处理数值特征
    processed_data = preprocessor.transform(input_df[numeric_features])

    # 预处理分类特征
    for col in categorical_features:
        le = label_encoders[col]
        # 处理未知类别
        if input_df[col].iloc[0] not in le.classes_:
            most_frequent = le.classes_[0]  # 使用第一个类别作为默认
            input_df[col] = most_frequent
        encoded_col = le.transform(input_df[col])
        processed_data = np.column_stack((processed_data, encoded_col))

    # 特征选择
    selected_data = selector.transform(processed_data)

    # 预测
    proba = model.predict_proba(selected_data)[0, 1]
    prediction = model.predict(selected_data)[0]

    return {'欺诈概率': float(proba), '预测结果': int(prediction)}


def batch_predict(test_data, policy_ids):
    """
    批量预测测试数据

    参数:
        test_data (array): 预处理后的测试数据
        policy_ids (Series): 测试数据的保单ID

    返回:
        DataFrame: 包含保单ID的预测结果
    """
    # 加载模型
    model = joblib.load('best_model.pkl')

    # 预测
    test_pred_proba = model.predict_proba(test_data)[:, 1]
    test_pred = model.predict(test_data)

    # 创建结果DataFrame
    results_df = pd.DataFrame({
        'policy_id': policy_ids,
        'fraud_probability': test_pred_proba,
        'fraud_prediction': test_pred
    })

    return results_df